import os
import xml.etree.ElementTree as ET
from xml.dom import minidom
import sys

def generate_annotations_from_filenames():
    """直接从文件名生成箭头方向标注文件"""
    
    # 获取当前脚本所在目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    parent_dir = os.path.dirname(current_dir)
    
    print(f"当前脚本目录: {current_dir}")
    print(f"父目录: {parent_dir}")
    
    # 只处理箭头相关的文件夹
    folders = {
        'arrows': 'arrow',
        'variations': 'arrow_variation'
    }
    
    total_xml_count = 0
    
    for folder, pattern_type in folders.items():
        folder_path = os.path.join(parent_dir, folder)
        print(f"\n检查文件夹: {folder_path}")
        
        if not os.path.exists(folder_path):
            print(f"  ✗ 文件夹不存在: {folder_path}")
            continue
            
        # 统计 PNG 文件数量
        png_files = [f for f in os.listdir(folder_path) if f.endswith('.png')]
        print(f"  ✓ 找到 {len(png_files)} 个 PNG 文件")
        
        if len(png_files) == 0:
            print(f"  ! 警告: {folder} 中没有 PNG 文件")
            continue
            
        xml_count = 0
        for filename in png_files:
            # 解析文件名获取箭头方向信息
            try:
                if pattern_type == 'arrow':
                    # up_red_000.png → 类别: up_arrow
                    parts = filename.split('_')
                    if len(parts) >= 2:
                        direction = parts[0]  # up, down, left, right
                        class_name = f"{direction}_arrow"
                    else:
                        class_name = 'unknown_arrow'
                    # 箭头大致在中央区域
                    bbox = [70, 70, 154, 154]
                    
                elif pattern_type == 'arrow_variation':
                    # arrow_up_red_00.png → 类别: up_arrow
                    parts = filename.split('_')
                    if len(parts) >= 3 and parts[0] == 'arrow':
                        direction = parts[1]  # up, down, left, right
                        class_name = f"{direction}_arrow"
                    else:
                        # 尝试从其他格式解析
                        for direction in ['up', 'down', 'left', 'right']:
                            if direction in filename:
                                class_name = f"{direction}_arrow"
                                break
                        else:
                            class_name = 'unknown_arrow'
                    # 变体箭头的边界框
                    bbox = [60, 60, 164, 164]
                
                # 只标注方向，忽略颜色
                # 生成VOC格式标注
                success = create_voc_annotation(folder_path, filename, class_name, bbox)
                if success:
                    xml_count += 1
                    total_xml_count += 1
                    
            except Exception as e:
                print(f"  ✗ 处理文件 {filename} 时出错: {e}")
        
        print(f"  ✓ 在 {folder} 中生成 {xml_count} 个 XML 文件")
    
    print(f"\n总结: 共生成 {total_xml_count} 个箭头方向标注文件")
    
    # 验证生成的文件
    print("\n验证生成的 XML 文件:")
    for folder in folders.keys():
        folder_path = os.path.join(parent_dir, folder)
        if os.path.exists(folder_path):
            xml_files = [f for f in os.listdir(folder_path) if f.endswith('.xml')]
            print(f"  {folder}: {len(xml_files)} 个 XML 文件")
            
            # 显示前几个文件的标注类别
            if xml_files:
                print("  示例标注:")
                for xml_file in xml_files[:3]:  # 显示前3个
                    xml_path = os.path.join(folder_path, xml_file)
                    try:
                        tree = ET.parse(xml_path)
                        root = tree.getroot()
                        for obj in root.findall('object'):
                            name = obj.find('name').text
                            print(f"    {xml_file} -> {name}")
                    except:
                        pass

def create_voc_annotation(folder_path, filename, class_name, bbox):
    """创建VOC格式标注文件"""
    
    try:
        root = ET.Element("annotation")
        
        # 基础信息
        ET.SubElement(root, "folder").text = os.path.basename(folder_path)
        ET.SubElement(root, "filename").text = filename
        
        # 图片尺寸
        size = ET.SubElement(root, "size")
        ET.SubElement(size, "width").text = "224"
        ET.SubElement(size, "height").text = "224" 
        ET.SubElement(size, "depth").text = "3"
        
        # 目标物体 - 只标注箭头方向
        object_elem = ET.SubElement(root, "object")
        ET.SubElement(object_elem, "name").text = class_name
        ET.SubElement(object_elem, "pose").text = "Unspecified"
        ET.SubElement(object_elem, "truncated").text = "0"
        ET.SubElement(object_elem, "difficult").text = "0"
        
        # 边界框
        bndbox = ET.SubElement(object_elem, "bndbox")
        ET.SubElement(bndbox, "xmin").text = str(bbox[0])
        ET.SubElement(bndbox, "ymin").text = str(bbox[1])
        ET.SubElement(bndbox, "xmax").text = str(bbox[2])
        ET.SubElement(bndbox, "ymax").text = str(bbox[3])
        
        # 保存XML文件
        xml_filename = filename.replace('.png', '.xml')
        xml_path = os.path.join(folder_path, xml_filename)
        
        # 美化输出
        rough_string = ET.tostring(root, 'utf-8')
        reparsed = minidom.parseString(rough_string)
        
        with open(xml_path, 'w', encoding='utf-8') as f:
            f.write(reparsed.toprettyxml(indent="  "))
        
        return True
        
    except Exception as e:
        print(f"    ✗ 创建标注失败 {filename}: {e}")
        return False

if __name__ == "__main__":
    print("开始自动标注箭头方向...")
    generate_annotations_from_filenames()
    print("箭头方向自动标注完成！")